import torch
from torch import nn

class DynamicWindowQFormer(nn.Module):
    def __init__(self, input_dim, hidden_dim, num_heads, dropout=0.1):
        super(DynamicWindowQFormer, self).__init__()
        # 多头注意力机制处理输入
        self.multihead_attn = nn.MultiheadAttention(input_dim, num_heads)
        # 前馈神经网络进一步处理特征
        self.ffn = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Dropout(dropout),
            nn.Linear(hidden_dim, input_dim)
        )
        # 层归一化提升训练稳定性
        self.layer_norm1 = nn.LayerNorm(input_dim)
        self.layer_norm2 = nn.LayerNorm(input_dim)

    def forward(self, x, attention_mask):
        # 多头注意力机制计算，应用注意力掩码处理变长输入
        attn_output, _ = self.multihead_attn(x, x, x, key_padding_mask=attention_mask)
        x = self.layer_norm1(x + attn_output)
        # 前馈神经网络处理
        ffn_output = self.ffn(x)
        x = self.layer_norm2(x + ffn_output)
        return x

# 示例用法
input_dim = 512
hidden_dim = 2048
num_heads = 8
dropout = 0.1
qformer = DynamicWindowQFormer(input_dim, hidden_dim, num_heads, dropout)

# 模拟输入数据，x 为特征矩阵，attention_mask 为注意力掩码
x = torch.randn(10, 32, input_dim)
attention_mask = torch.ones(10, 32).bool()
output = qformer(x, attention_mask)  